# predict.py
import joblib
import numpy as np
# predict.py
from ai_transaction_detector import AITransactionDetector as LocalPredictor
from qwen_advisor import generate_optimization_advice


class TransactionPredictor:
    def __init__(self, model_path="models/anomaly_model.pkl"):
        self.model = joblib.load(model_path)
        self.scaler = joblib.load(model_path.replace(".pkl", "_scaler.pkl"))

    def predict(self, tx):
        X = np.array([[tx['duration_sec'], tx['wait_time_ms'], tx['rows_affected'], tx['lock_count']]])
        X_scaled = self.scaler.transform(X)

        pred = self.model.predict(X_scaled)[0]  # 1: 正常, -1: 异常
        score = self.model.decision_function(X_scaled)[0]

        is_anomaly = pred == -1
        return {
            "is_long_running": is_anomaly,
            "confidence": abs(score),
            "risk_level": "高" if is_anomaly else "低",
            "suggestion": "建议监控或回滚" if is_anomaly else "运行正常"
        }

class AITransactionAdvisor:
    def __init__(self):
        self.local_model = LocalPredictor()  # 本地 IsolationForest

    def predict(self, tx, sql="", explain=""):
        # 第一步：本地模型判断是否高风险
        local_result = self.local_model.predict(tx)

        if not local_result["is_long_running"]:
            return {
                "is_long_running": False,
                "risk_score": local_result["risk_score"],
                "suggestion": "事务正常，无需优化"
            }

        # 第二步：调用通义千问生成专业建议
        advice = generate_optimization_advice(tx, sql=sql, explain_output=explain)

        return {
            "is_long_running": True,
            "risk_score": local_result["risk_score"],
            "suggestion": advice
        }

# 演示
if __name__ == "__main__":
    pred = TransactionPredictor()

    # 测试正常事务
    normal_tx = {"duration_sec": 10, "wait_time_ms": 20, "rows_affected": 100, "lock_count": 1}
    print("🟢 正常事务检测:", pred.predict(normal_tx))

    # 测试异常事务
    long_tx = {"duration_sec": 500, "wait_time_ms": 1000, "rows_affected": 5, "lock_count": 7}
    print("🔴 长事务检测:", pred.predict(long_tx))